Description

Makes use_standalone_converter work end to end for Qwen3.5 against current tpu-inference, so an RL run can roll out on vLLM's native Qwen3.5 (torchax) implementation with weight sync from MaxText. Intended to unblock a validation run of the deepswe e2e (google/tunix abb9988) on the native-vLLM sampler; the longer-term canonical-layout path is #5034 / vllm-project/tpu-inference#3477 / google/tunix#2014, which moves the layout knowledge back into tpu-inference.

1. torchax_converter/qwen35_moe.py: match the current tpu-inference runner layout

The converter's output had drifted from the runner state (the symptom behind #4713: shape mismatches and gibberish). Fixes, all verified tensor-by-tensor against a runner-state dump of Qwen/Qwen3.5-35B-A3B loaded from the HF checkpoint on current tpu-inference main:

• Fused projections (qkv_proj, in_proj_qkvz, in_proj_ba, shared_expert.gate_up_proj) get a final transpose: the runner stores linear kernels [in, out] with the out dimension reordered per TP shard.
• o_proj, GDN out_proj and shared_expert.down_proj are no longer transposed: the MaxText kernels ([H*dh, D], [F, D]) are already the runner's [in, out] orientation.
• KV heads are replicated (each head tp // num_kv_heads times, consecutively) before the per-shard interleave when tp exceeds the KV head count, mirroring vLLM's QKVParallelLinear (Qwen3.5-35B has 2 KV heads, so this triggers at TP ≥ 4).
• Routed-expert tensors are additionally emitted under the current mlp.experts.routed_experts.{w13,w2}_weight names (same arrays, no copy); the old mlp.experts.* names are kept for older vLLM builds, and the sync drops whichever the target lacks.
• A_log stays float32 (the runner keeps it in float32; bfloat16 would round it).

2. maxtext_vllm_rollout.py: actually run the standalone converter

use_standalone_converter existed in RLConfig but was only read by validate_converter.py, and MaxTextVllmSampler stored the converter without ever applying it. Now:

• _create_model_converter instantiates the Qwen3.5 (or Gemma4) standalone converter when use_standalone_converter is set (config field or vllm_additional_config key).
• MaxTextVllmSampler.update_params routes standalone converters through a new _sync_standalone_converted: KV-cache teardown → convert() → chunked, Pathways-aware reshard of each tensor onto its existing runner sharding (reshard_chunk_size respected; falls back to a single reshard on older Tunix) → in-place assignment into the runner's flat state dict → state_leaves refresh → KV-cache rebuild. Shape/dtype mismatches raise with a message pointing at converter/tpu-inference drift, and uncovered runner tensors are logged loudly.

3. Expert parallelism and attention DP

The converter is parameterized on the sampler's sharding (vllm_attn_dp, vllm_use_ep), with the rules derived empirically from runner-state dumps under the production configuration (TP=8, --enable-expert-parallel with USE_MOE_EP_KERNEL=0 → GMM_EP, enable_dp_attention + attn_dp_size=4, production env/LIBTPU flags):

• attention, GDN and the shared expert are sharded tp // attn_dp ways — their interleave and the KV-head-replication threshold use that count (at attn_dp=4 with 2 KV heads there is no replication; qkv_proj is 9216 wide, not 12288);
• under expert parallelism the routed experts take the GMM_EP layout: sharded on the expert axis, w13 = [E, D, 2·pad128(F)] with no per-TP gate/up interleave.

MaxTextVllmRollout fills the hints automatically: attn_dp_size from additional_config.sharding.sharding_strategy (when enable_dp_attention is set) and enable_expert_parallel from the vLLM engine kwargs. Defaults reproduce plain TP bit-for-bit.

How to use in the e2e

In the sampler config: use_standalone_converter=True, and make vLLM run its native model — remove vllm_hf_overrides={"architectures": ["MaxTextForCausalLM"]} and the maxtext_config entry from rollout_vllm_additional_config (keep MODEL_IMPL_TYPE=vllm / NEW_MODEL_DESIGN=1).

Caveats

This file intentionally mirrors tpu-inference's internal layout: it must be revalidated whenever tpu-inference changes its weight processing, and it does not cover quantized (FP8) checkpoints — a bf16-trainer → fp8-sampler sync needs online quantization, which no path has yet. The canonical-layout sync (#5034 + tpu-inference#3477 + tunix#2014) removes the layout coupling.

Tests

• Accuracy: gsm8k via lm_eval --model local-chat-completions --tasks gsm8k --gen_kwargs max_gen_toks=2048 --apply_chat_template --fewshot_as_multiturn against vllm serve Qwen/Qwen3.5-35B-A3B under the same sharding/env (TP=8, --enable-expert-parallel, attn_dp_size=4, tuned serving flags): 0.9621 flexible / 0.9629 strict. The synced state was proven bit-identical to this HF-loaded state, so the score applies to the synced model.

• Live e2e on TPU (v7x-8): MaxText 35B weights (HF-converted checkpoint, restacked to the scanned layout) → MaxTextVllmSampler with Qwen35MaxTextToVLLMConverter → _sync_standalone_converted into a running native vLLM Qwen3.5 engine at TP=8 + attn_dp_size=4 + --enable-expert-parallel whose weights had been overwritten with random values (generating confirmed garbage). After the sync: every dumped runner tensor bit-identical to the HF-loaded reference, and 4/4 greedy generations identical to the in-process HF-weight reference. Sync wall time 36 s (host-side convert in this harness; in the real run the trainer state is already on TPU).

• CPU validation: unscanned gs://maxtext-model-checkpoints/qwen3.5-35b-a3b/unscanned/0/items restacked into the scanned block layout → Qwen35MaxTextToVLLMConverter.convert() → every emitted tensor compared bit-wise against runner-state dumps of the same model loaded from the HF checkpoint by tpu-inference itself (v7x-8, current main), in both sharding modes:

  • plain TP=8: all 573 runner weight tensors identical (full-model dump);
  • TP=8 + attn_dp_size=4 + --enable-expert-parallel (USE_MOE_EP_KERNEL=0 → GMM_EP) with the production env/LIBTPU flags: all compared tensors identical (layers 0–7 + globals — two full 4-slot cycles).

• Coverage check confirms every runner weight is produced (the extra emitted keys are the old/new expert-name aliases).

• python -m py_compile on both files.